//----------------------------------------------------------------------------
// M A R I T I M E  R E S E A R C H  I N S T I T U T E  N E T H E R L A N D S
//----------------------------------------------------------------------------
// Copyright (C) 2011 - MARIN - All rights reserved - http://www.marin.nl
//----------------------------------------------------------------------------
// Program    : mermaid
// Module     : vbm
// File       : cg_kernels.h
// Author     : M.A. de Jong
//----------------------------------------------------------------------------

#ifndef CG_KERNELS_H
#define	CG_KERNELS_H

#define Bx blockDim.x
#define bx blockIdx.x
#define tx threadIdx.x
#define By blockDim.y
#define by blockIdx.y
#define ty threadIdx.y



namespace cg
{

// AXPY kernel variant 1: y = a*x + y
template <class T>
__global__ void AXPY(T *y,          
                     const T *x,
                     const T alpha, 
                     const T beta,
                     const int n)
{
    int tid = bx * Bx + tx;
    if (tid < n)
        y[tid] = alpha * x[tid] + beta * y[tid];
}

// Makes diagonal preconditioning matrix P
template <class T>
__global__ void makeDiagonal(T *stP, const T *stC, const int n)
{
    int tid;

    tid = bx * Bx + tx;
    if (tid < n)
        stP[tid] = sqrt(stC[tid]);
}

// Makes IP preconditioner on device
template <class T>
__global__ void makeIPPrec(T *stKC, T *stKS, T *stKW,
                           const T *stAC, const T *stAS, const T *stAW,
                           const int Nx1, const int Nx2)
{
    int tid = bx * Bx + tx;

    T val;
    T valKW;
    T valKS;
    T valKC;

    if (tid < Nx1 * Nx2)
    {
        val = 1.0 / stAC[tid];
        valKW = -stAW[tid] * val;
        valKS = -stAS[tid] * val;

        if (tid == 0)
            valKC = 1.0;
        else if (tid > 0 && tid < Nx1)
            valKC = 1.0 + valKW * valKW;
        else
            valKC = 1.0 + valKW * valKW + valKS * valKS;

        stKW[tid] = valKW;
        stKS[tid] = valKS;
        stKC[tid] = valKC;
    }
}



// Updates the stencils of matrix A
template <class T>
__global__ void updateMatrix(T *stC, T *stS, T *stW, 
                        const T *stP, const int Nx1, const int Nx2)
{
    int tid = bx * Bx + tx;
    if (tid < Nx1 * Nx2) {
        stC[tid] = 1.0; // center (not needed explicitly)
        if (tid > 0)
            stW[tid] /= (stP[tid - 1  ] * stP[tid]); // west
        if (tid > Nx1 - 1)
            stS[tid] /= (stP[tid - Nx1] * stP[tid]); // south
    }
}

// Element-wise multiply; required for x = P^T x
template <class T>
__global__ void elementWiseMul(T *x, const T *p, const int n)
{
    int tid = bx * Bx + tx;
    if (tid < n)
        x[tid] *= p[tid];
}


// Element-wise divide; required for x = P^{-T} x
template <class T>
__global__ void elementWiseDiv(T *x, const T *p, const int n)
{
    int tid = bx * Bx + tx;
    if (tid < n)
        x[tid] /= p[tid];
}

// Computes y = Ax for sparse A, matrix A is given by matc, mats, matw
// using all 4 stencils (and no textures for them) is fastest!
template <class T>
__global__ void SpMVv1(T *y,
                     const T *stC,    
                     const T *stS,
                     const T *stW,
                     const T *x,
                     const int Nx1, 
                     const int Nx2)
{
    int tid = bx * Bx + tx;

    T sum = 0;

    if (tid < Nx1 * Nx2)
    {
        sum = stC[tid] * fetchX(x, tid);                  // center

        if (tid + Nx1 < Nx1 * Nx2)
            sum += stS[tid + Nx1] * fetchX(x, tid + Nx1); // north
        if (tid + 1 < Nx1 * Nx2)
            sum += stW[tid + 1]   * fetchX(x, tid + 1);   // east
        if (tid - Nx1 >= 0)
            sum += stS[tid]       * fetchX(x, tid - Nx1); // south
        if (tid - 1 >= 0)
            sum += stW[tid]       * fetchX(x, tid - 1);   // west
 
        y[tid] = sum;
    }
}

// Computes y = Ax for sparse A, matrix A is given by matc, mats, matw
// NOTE: diagonally scaled matrix A, so 1's on main diagonal
template <class T>
__global__ void SpMVv2(T *y,
                     const T *stC, // note: not used (because all 1's)  
                     const T *stS,
                     const T *stW,
                     const T *x,
                     const int Nx1, 
                     const int Nx2)
{
    int tid = bx * Bx + tx;

    T sum = 0;

    if (tid < Nx1 * Nx2)
    {
        sum = fetchX(x, tid);                             // center

        if (tid + Nx1 < Nx1 * Nx2)
            sum += stS[tid + Nx1] * fetchX(x, tid + Nx1); // north
        if (tid + 1 < Nx1 * Nx2)
            sum += stW[tid + 1]   * fetchX(x, tid + 1);   // east
        if (tid - Nx1 >= 0)
            sum += stS[tid]       * fetchX(x, tid - Nx1); // south
        if (tid - 1 >= 0)
            sum += stW[tid]       * fetchX(x, tid - 1);   // west

        y[tid] = sum;
    }
}

} // namespace

#endif // CG_KERNELS_H

